# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import random
import numpy as np
import torch

from fairseq.data import data_utils, FairseqDataset


def collate(samples, pad_idx, eos_idx, sep_idx=None, mid_idx=None, use_context=None, mask_mid=False, mask_input=False):
    if len(samples) == 0:
        return {}

    def merge(key, _pad_idx=pad_idx, is_list=False):
        if is_list:
            res = []
            for i in range(len(samples[0][key])):
                res.append(data_utils.collate_tokens(
                    [s[key][i] for s in samples], _pad_idx, eos_idx, left_pad=False,
                ))
            return res
        else:
            return data_utils.collate_tokens(
                [s[key] for s in samples], _pad_idx, eos_idx, left_pad=False,
            )

    src_tokens = merge('source')
    context = None
    context_length = None
    tgt_prompt_mask = None
    tgt_entity = None
    if samples[0]['target'] is not None:
        is_target_list = isinstance(samples[0]['target'], list)
        target = merge('target', is_target_list)
        target_lengths = [len(s['target']) for s in samples]

        if samples[0].get('target_label', None) is not None:
            # event label and entity label
            tgt_entity = merge('target_label', _pad_idx=-1)

            assert target.size() == tgt_entity.size()

        # for pad == eos in bpe, so mask by length
        tgt_prompt_mask = torch.ones_like(target).type(torch.uint8)
        for i in range(len(samples)):
            tgt_prompt_mask[i][target_lengths[i]:] = 0

        if mask_mid and mid_idx is not None and sep_idx is not None:
            def get_segment(item):
                start, end = 0, 0
                for i in range(len(item)):
                    if item[i: i + len(mid_idx)].tolist() == mid_idx:
                        start = i
                    if item[i: i + len(sep_idx)].tolist() == sep_idx:
                        end = i
                        break
                if end < start:
                    end = start + int((item.size(0) - start) * 0.15)
                return start, end

            for i in range(len(samples)):
                start, end = get_segment(samples[i]['target'])
                tgt_prompt_mask[i][start:end] = 0

                if mask_input:
                    mask_percent = 0.15
                    mask_len = int((end - start) * mask_percent)
                    mask_index = random.sample(range(start+1, end+1), k=mask_len)
                    src_tokens[i][mask_index] = 50009

        if use_context and sep_idx is not None:
            def get_len(item):
                for i in range(len(item)):
                    if item[i: i + len(sep_idx)].tolist() == sep_idx:
                        return i + len(sep_idx)
                return len(item)

            context_length = list(map(get_len, [s['target'] for s in samples]))
            max_len = max(context_length)
            context = target.new(target.size(0), max_len).long().fill_(pad_idx)
            for i in range(len(samples)):
                context[i, :context_length[i]] = target[i, :context_length[i]]
    else:
        target = src_tokens

    batch = {
        'id': torch.LongTensor([s['id'] for s in samples]),
        'nsentences': len(samples),
        'ntokens': tgt_prompt_mask.nonzero().size(0),
        'net_input': {
            'src_tokens': src_tokens,
            'src_lengths': torch.LongTensor([
                s['source'].numel() for s in samples
            ]),
        },
        'target': target,
        'tgt_prompt_mask': tgt_prompt_mask,
        'tgt_entity': tgt_entity,
    }
    if context is not None:
        batch['context'] = context
        batch['net_input']['context_lengths'] = torch.LongTensor(context_length)
    return batch


class MonolingualCorefDataset(FairseqDataset):
    """
    A wrapper around torch.utils.data.Dataset for monolingual data.

    Args:
        dataset (torch.utils.data.Dataset): dataset to wrap
        sizes (List[int]): sentence lengths
        vocab (~fairseq.data.Dictionary): vocabulary
        shuffle (bool, optional): shuffle the elements before batching
            (default: True).
    """

    def __init__(self, dataset, sizes, src_vocab, tgt_vocab, add_eos_for_other_targets, shuffle,
                 targets=None, add_bos_token=False, use_context=False, truncate=False, max_sample_len=1024,
                 mask_mid=False, mask_input=False, target_labels=None):
        self.dataset = dataset
        self.sizes = np.array(sizes)
        self.vocab = src_vocab
        self.tgt_vocab = tgt_vocab
        self.add_eos_for_other_targets = add_eos_for_other_targets
        self.shuffle = shuffle
        self.add_bos_token = add_bos_token

        self.use_context = use_context
        self.truncate = truncate
        self.max_sample_len = max_sample_len
        self.mask_mid = mask_mid
        self.mask_input = mask_input
        self.target_labels = target_labels

        if self.truncate:
            self.sizes[self.sizes > max_sample_len] = max_sample_len

        assert targets is None or all(t in {'self', 'future', 'past'} for t in targets), \
            "targets must be none or one of 'self', 'future', 'past'"
        if targets is not None and len(targets) == 0:
            targets = None
        self.targets = targets

        if self.vocab.index('<SEP>') != self.vocab.unk():
            self.sep_idx = self.vocab.index('<SEP>')
        elif getattr(self.vocab, 'enc', False):
            self.sep_idx = self.vocab.enc.encode(' <SEP>')
        else:
            self.sep_idx = None
        if self.vocab.index('<S>') != self.vocab.unk():
            self.mid_idx = self.vocab.index('<S>')
        elif getattr(self.vocab, 'enc', False):
            self.mid_idx = self.vocab.enc.encode(' <S>')
        else:
            self.mid_idx = None

    def __getitem__(self, index):
        if self.targets is not None:
            # *future_target* is the original sentence
            # *source* is shifted right by 1 (maybe left-padded with eos)
            # *past_target* is shifted right by 2 (left-padded as needed)
            #
            # Left-to-right language models should condition on *source* and
            # predict *future_target*.
            # Right-to-left language models should condition on *source* and
            # predict *past_target*.
            source, future_target, past_target = self.dataset[index]
            source, target = self._make_source_target(source, future_target, past_target)

            target_label = self.target_labels[index] if self.target_labels is not None else None
        else:
            source = self.dataset[index]
            target = None
            target_label = None
        source, target, target_label = self._maybe_add_bos(source, target, target_label)

        if self.truncate:
            source = source[:self.max_sample_len]
            if target is not None:
                target = target[:self.max_sample_len]
                target_label = target_label[:self.max_sample_len]
                assert target.size() == target_label.size()
        return {'id': index, 'source': source, 'target': target, 'target_label': target_label}

    def __len__(self):
        return len(self.dataset)

    def _make_source_target(self, source, future_target, past_target):
        if self.targets is not None:
            target = []

            if self.add_eos_for_other_targets and (('self' in self.targets) or ('past' in self.targets)) \
                    and source[-1] != self.vocab.eos():
                # append eos at the end of source
                source = torch.cat([source, source.new([self.vocab.eos()])])

                if 'future' in self.targets:
                    future_target = torch.cat([future_target, future_target.new([self.vocab.pad()])])
                if 'past' in self.targets:
                    # first token is before the start of sentence which is only used in "none" break mode when
                    # add_eos_for_other_targets is False
                    past_target = torch.cat([past_target.new([self.vocab.pad()]), past_target[1:], source[-2, None]])

            for t in self.targets:
                if t == 'self':
                    target.append(source)
                elif t == 'future':
                    target.append(future_target)
                elif t == 'past':
                    target.append(past_target)
                else:
                    raise Exception('invalid target ' + t)

            if len(target) == 1:
                target = target[0]
        else:
            target = future_target

        return source, self._filter_vocab(target)

    def _maybe_add_bos(self, source, target, target_label):
        if self.add_bos_token:
            source = torch.cat([source.new([self.vocab.bos()]), source])
            if target is not None:
                target = torch.cat([target.new([self.tgt_vocab.bos()]), target])
                target_label = torch.cat([target_label.new([0]), target_label])
        return source, target, target_label

    def _filter_vocab(self, target):
        if len(self.tgt_vocab) != len(self.vocab):
            def _filter(target):
                mask = target.ge(len(self.tgt_vocab))
                if mask.any():
                    target[mask] = self.tgt_vocab.unk()
                return target

            if isinstance(target, list):
                return [_filter(t) for t in target]
            return _filter(target)
        return target

    def collater(self, samples):
        """Merge a list of samples to form a mini-batch.

        Args:
            samples (List[dict]): samples to collate

        Returns:
            dict: a mini-batch with the following keys:

                - `id` (LongTensor): example IDs in the original input order
                - `ntokens` (int): total number of tokens in the batch
                - `net_input` (dict): the input to the Model, containing keys:

                  - `src_tokens` (LongTensor): a padded 2D Tensor of tokens in
                    the source sentence of shape `(bsz, src_len)`. Padding will
                    appear on the right.

                - `target` (LongTensor): a padded 2D Tensor of tokens in the
                  target sentence of shape `(bsz, tgt_len)`. Padding will appear
                  on the right.
        """
        return collate(samples, self.vocab.pad(), self.vocab.eos(),
                       sep_idx=self.sep_idx, mid_idx=self.mid_idx,
                       use_context=self.use_context, mask_mid=self.mask_mid,
                       mask_input=self.mask_input)

    def num_tokens(self, index):
        """Return the number of tokens in a sample. This value is used to
        enforce ``--max-tokens`` during batching."""
        return self.sizes[index]

    def size(self, index):
        """Return an example's size as a float or tuple. This value is used when
        filtering a dataset with ``--max-positions``."""
        return self.sizes[index]

    def ordered_indices(self):
        """Return an ordered list of indices. Batches will be constructed based
        on this order."""
        if self.shuffle:
            order = [np.random.permutation(len(self))]
        else:
            order = [np.arange(len(self))]
        order.append(self.sizes)
        return np.lexsort(order)

    @property
    def supports_prefetch(self):
        return getattr(self.dataset, 'supports_prefetch', False)

    def prefetch(self, indices):
        self.dataset.prefetch(indices)
        if self.target_labels is not None:
            self.target_labels.prefetch(indices)
